from langChain import hub
from langChain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS, Chroma
import os

from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.tools import create_retriever_tool
from langchain_text_splitters import RecursiveCharacterTextSplitter

from langChain.config import model
from langChain.config import embedding
from langchain_tavily import TavilySearch

os.environ["TAVILY_API_KEY"] = "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"

# search = TavilySearch(max_results=1)
# print(search.invoke("合肥天气怎么样?"))

# 拆文档下载、拆分、存储
# 1.文档下载
loader = WebBaseLoader("https://zh.wikipedia.org/wiki/%E7%8C%AB")
docs = loader.load()

# 2.大文本切割
documents = RecursiveCharacterTextSplitter(
    # chunk_size 参数在 RecursiveCharacterTextSplitter 中用于指定每个文档块的最大字符数。它的作用主要有以下几个方面：
    # chunk_overlap 参数用于指定每个文档块之间的重叠字符数。这意味着，当文档被拆分成较小的块时，每个块的末尾部分会与下一个块的开头部分有一定数量的重叠字符。
    # 第一个块包含字符 1 到 1000。第二个块包含字符 801 到 1800。第三个块包含字符 1601 到 2600。
    chunk_size=1000, chunk_overlap=200
).split_documents(docs)

# 3.保存到向量数据库
vector = FAISS.from_documents(documents, embedding=embedding)
# vector = Chroma.from_documents(documents, embedding)

# 4.获取检索器
retriever = vector.as_retriever()

# 检索内容
# print('检索的内容 ：：', retriever.invoke("猫的特征")[0])

# 创建工具
search = TavilySearch(max_results=1)
retriever_tool = create_retriever_tool(retriever, "wiki_search", "搜索维基百科", )
tools = [search, retriever_tool]
model_with_tools = model.bind_tools(tools)
# response = model_with_tools.invoke([HumanMessage(content="今天上海天气怎么样")])
# print(f"ContentString: {response.content}")
# print(f"ToolCalls: {response.tool_calls}")

# 获取要使用的提示 - 您可以修改这个！
# 如果您想查看此提示的内容并访问LangSmith，
# 您可以转到:https://smith.langchain.com/hub/hwchase17/openai-functions-agent
prompt = hub.pull("hwchase17/openai-functions-agent")
# 返回的是一个 Agent 对象，具备调用工具的能力，但尚未与完整的执行器框架集成。
agent = create_tool_calling_agent(model, tools, prompt)
# 返回的是一个 AgentExecutor 对象，提供更完整的执行流程控制，内部已经绑定了 Agent 和工具集。
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# print(agent_executor.invoke({"input": "你好"}))
# print(agent_executor.invoke({"input": "猫的特征"}))
# print('agent调用结果：',agent.invoke({'messages': [HumanMessage(content="猫的特征?今天上海天气怎么样?")]}))  # 错误
# print('agent调用结果：',agent_executor.invoke({"input": "猫的特征?今天上海天气怎么样?"}))

# 添加历史记录
# # 简单示例
# # print(agent_executor.invoke({"input": "你好，我的名字是Cyber", "chat_history": []}))
#
# response = agent_executor.invoke({
#     "chat_history": [HumanMessage(content="Hi，我的名字是Jack"),
#                      AIMessage(content="你好，Jack，很高兴见到你！有什么我可以帮助你的吗？"),
#                      ],
#     "input": "我的名字是什么?",
# })
# print('添加历史记录以后的回答 ：', response)

# 生产示例
store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]


agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

response = agent_with_chat_history.invoke(
    {"input": "Hi，我的名字是Jack"},
    config={"configurable": {"session_id": "123"}},
)

print('第一次 ：', response)

response = agent_with_chat_history.invoke(
    {"input": "我叫什么名字?"},
    config={"configurable": {"session_id": "123"}},
)

print('第二次 ：', response)

response = agent_with_chat_history.invoke(
    {"input": "我叫什么名字?"},
    config={"configurable": {"session_id": "456"}},
)

print('第三次 ,更换了session_id：', response)
